from .elastic_manager import es


def get_search_suggest(target, keyword):
    search_body = {
        "suggest": {
            "getSuggest": {
                "prefix": keyword,
                "completion": {
                    "field": "suggest",
                    "size": 10
                }
            }
        }
    }
    res = es.search(index=target, body=search_body)
    hits = res['suggest']['getSuggest'][0]['options']
    suggestions = []
    for option in hits:
        suggestions.append(option['text'])
    return suggestions


def old_get_search_result(keyword):
    # 构建搜索查询体
    search_body = {
        "query": {
            "bool": {
                "should": [
                    # 共有字段
                    {"match": {"display_name": {"query": keyword, "boost": 100}}},
                    {"match": {"display_name_alternatives": {"query": keyword, "boost": 20}}},

                    # paper 索引字段
                    # {"match": {"title": {"query": keyword, "boost": 100}}},
                    {"match": {"keywords.keyword.keyword": {"query": keyword, "boost": 25}}},
                    # {"match": {"primary_location.source.display_name": {"query": keyword, "boost": 50}}},
                    # {"match": {"locations.source.display_name": {"query": keyword, "boost": 10}}},
                    {"match": {"authorships.author.display_name": {"query": keyword, "boost": 50}}},
                    # {"match": {"concepts.display_name": {"query": keyword, "boost": 0.2}}},
                    {"match": {"abstract": {"query": keyword, "boost": 40}}},
                    # 其他字段
                    # {"match": {"last_known_institution.display_name": {"query": keyword, "boost": 10}}},
                    {"match": {"x_concepts.display_name": {"query": keyword, "boost": 30}}},
                    # {"match": {"host_organization_name": {"query": keyword, "boost": 30}}},
                    # {"match": {"associated_institutions.display_name": {"query": keyword, "boost": 10}}},
                    # {"match": {"international.display_name.*": {"query": keyword, "boost": 20}}},
                    {"match": {"description": {"query": keyword, "boost": 30}}},
                    # {"match": {"international.description.*": {"query": keyword, "boost": 10}}},
                    # {"match": {"ancestors.display_name": {"query": keyword, "boost": 10}}},
                    # {"match": {"related_concepts.display_name": {"query": keyword, "boost": 2}}}
                ]
            }
        },
        "size": 500,  # 指定返回结果数量
        "_source": [
            # 共有字段
            "id",
            "type",
            "display_name",
            # "display_name_alternatives",

            # paper 索引字段
            # "doi",
            "title",
            "language",

            # "primary_location.pdf_url",
            # "primary_location.landing_page_url",
            "primary_location.source.id",
            "primary_location.source.display_name",

            "primary_location.pdf_url",
            # "locations.landing_page_url",
            # "locations.source.id",
            # "locations.source.display_name",

            # "authorships.author_position",
            "authorships.author.id",
            "authorships.author.display_name",
            # "keywords",
            "concepts.id",
            "concepts.display_name",
            # "concepts.score",
            # "concepts.level",
            # "abstract",

            # scholar venue org 共有字段
            "works_count",
            "cited_by_count",

            "x_concepts.id",
            "x_concepts.display_name",
            # "x_concepts.score",
            # "x_concepts.level",

            # scholar 索引字段
            "last_known_institution.id",
            "last_known_institution.display_name",

            # venue org 共有字段
            "homepage_url",

            # venue 索引字段

            # org concept 共有字段
            "image_url",

            # org 索引字段
            "country_code",

            # concept 索引字段
            "level",
            "description",

            "ancestors.id",
            "ancestors.display_name",
            # "ancestors.level",

            # "related_concepts.id",
            # "related_concepts.display_name",
            # "related_concepts.score",
            # "related_concepts.level"

        ]
    }
    # print(search_body)
    indices = ['paper', 'scholar', 'venue', 'org', 'concept']
    result = es.search(index=indices, body=search_body)
    return result['hits']


def get_query_body(target, keyword, agg):
    should_body = []
    if target == 'paper':
        should_body = [
            {"match": {"display_name": {"query": keyword, "boost": 8}}},
            {"match": {"keywords.keyword": {"query": keyword, "boost": 4}}},
            {"match": {"abstract": {"query": keyword, "boost": 3}}},
            {"match": {"concepts.display_name": {"query": keyword, "boost": 4}}}
        ]
    elif target == 'scholar':
        should_body = [
            {"match": {"display_name": {"query": keyword, "boost": 6}}},
            {"match": {"display_name_alternatives": {"query": keyword, "boost": 6}}},
            {"match": {"x_concepts.display_name": {"query": keyword, "boost": 6}}},
            {"match": {"last_known_institution.display_name": {"query": keyword, "boost": 2}}},
        ]
    elif target == 'venue':
        should_body = [
            {"match": {"display_name": {"query": keyword, "boost": 7}}},
            {"match": {"x_concepts.display_name": {"query": keyword, "boost": 3}}},
        ]
    elif target == 'concept':
        should_body = [
            {"match": {"display_name": {"query": keyword, "boost": 7}}},
            {"match": {"description": {"query": keyword, "boost": 3}}}
        ]
    elif target == 'org':
        should_body = [
            {"match": {"display_name": {"query": keyword, "boost": 8}}},
            {"match": {"display_name_alternatives": {"query": keyword, "boost": 4}}},
            {"match": {"international.display_name": {"query": keyword, "boost": 8}}},
            {"match": {"x_concepts.display_name": {"query": keyword, "boost": 6}}},
        ]
    must_body = [{"bool": {"should": should_body}}]
    if agg is not None and len(agg) >= 1:  # agg非空时添加新的约束
        if target != 'paper':
            type_constraints = agg.get('type', None)
            level_constraints = agg.get('level', None)
            key_array = []
            value_array = []
            if type_constraints is not None:
                for constraint in type_constraints:
                    key_array.append(get_agg_real_str_by_key('type'))
                    value_array.append(constraint)
            if level_constraints is not None:
                for constraint in level_constraints:
                    key_array.append(get_agg_real_str_by_key('level'))
                    value_array.append(constraint)
            i = 0
            buckets = []
            for key in key_array:
                temp_bool_body = {
                    "match": {
                        key: value_array[i]
                    }
                }
                buckets.append(temp_bool_body)
                i = i + 1
            must_body.append({"bool": {"should": buckets}})
        else:
            # print(agg)
            concept_constraints = agg.get('concepts', None)
            venue_constraints = agg.get('venue', None)
            type_constraints = agg.get('type', None)
            concept_should_json = create_should_json(concept_constraints, 'concepts.display_name.keyword')
            venue_should_json = create_should_json(venue_constraints, get_agg_real_str_by_key('venue'))
            type_should_json = create_should_json(type_constraints, 'type.keyword')
            all_jsons = []
            if concept_should_json is not None:
                all_jsons.append(concept_should_json)
            if venue_should_json is not None:
                all_jsons.append(venue_should_json)
            if type_should_json is not None:
                all_jsons.append(type_should_json)
            must_body.append({
                "bool": {
                    "must": all_jsons
                }
            })
    query_body = {
        "bool": {
            "must": must_body
        }
    }
    # print("paper:")
    # print(query_body)
    return query_body


def create_should_json(array, field):
    if array is None:
        return None
    should_matches = []
    for s in array:
        should_matches.append({
            "match": {
                field: s
            }
        })
    return {
        "bool": {
            "should": should_matches
        }
    }


def get_agg_real_str_by_key(key):
    """
    将简化字符串转化为用于ES search的属性名字符串
    """
    if key == 'concepts':
        return 'concepts.display_name.keyword'
    elif key == 'venue':
        return 'primary_location.source.display_name.keyword'
    elif key == 'level':
        return key
    else:
        return key + '.keyword'


def get_display_source(target):
    source = ["display_name", "cited_by_count"]
    if target == 'paper':
        source.append("language")
        source.append("authorships.author.id")
        source.append("authorships.author.display_name")
        source.append("concepts.id")
        source.append("concepts.display_name")
        source.append("publication_date")
        source.append("primary_location.pdf_url")
    elif target == 'scholar':
        source.append("last_known_institution.display_name")
        source.append("x_concepts.display_name")
        source.append("x_concepts.id")
        source.append("works_count")
    elif target == 'org':
        source.append("type")
        source.append("works_count")
        source.append("x_concepts.display_name")
        source.append("x_concepts.id")
    elif target == 'venue':
        source.append("works_count")
        source.append("x_concepts.display_name")
        source.append("x_concepts.id")
        source.append("issn_l")
        source.append("country_code")
    elif target == 'concept':
        source.append("description")
        source.append("works_count")
    return source


def get_search_count_pre_target(target, keyword, agg=None):
    search_body = {
        "query": get_query_body(target, keyword, agg)
    }
    resp = es.count(index=target, body=search_body)
    return resp['count']


def get_search_counts(keyword, agg):
    return {
        "paper_count": get_search_count_pre_target("paper", keyword, agg=agg),
        "scholar_count": get_search_count_pre_target("scholar", keyword),
        "venue_count": get_search_count_pre_target("venue", keyword),
        "org_count": get_search_count_pre_target("org", keyword),
        "concept_count": get_search_count_pre_target("concept", keyword)
    }


def get_search_result(target, res_from, res_size, keyword, agg=None, sort_key=None):
    query_body = get_query_body(target, keyword, agg)
    source_body = get_display_source(target)
    search_body = {
        "query": query_body,
        "from": res_from,
        "size": res_size,
        "_source": source_body
    }

    if sort_key is not None:
        sort_body = [{sort_key: {"order": "desc"}}]
        search_body['sort'] = sort_body
    # print(search_body)
    print(search_body, end="\n=======================")
    resp = es.search(index=target, body=search_body)
    return resp['hits']['hits']


def target_2_query_fields(target):
    res = []
    if target == 'paper':
        res = ["display_name", "concepts.display_name", "keywords.keyword", "abstract"]
    elif target == 'venue':
        res = ["display_name", "x_concepts.display_name"]
    elif target == 'org':
        res = ["display_name", "x_concepts.display_name", "display_name_alternatives", "international.display_name"]
    elif target == 'concept':
        res = ["display_name", "description"]
    return res


def buckets_search(target, field, keyword):
    # TODO : VENUE TYPE ORG, TYpe  paper primary_location.source.display_name
    key_str = get_agg_real_str_by_key(field)
    mult_fields = target_2_query_fields(target)
    search_body = {
        "query": {
            "bool": {
                "should": [
                    {
                        "multi_match": {
                            "query": keyword,
                            "fields": mult_fields
                        }
                    }
                ]
            }
        },
        "size": 0,
        "aggs": {
            "my_agg": {
                "terms": {
                    "field": key_str,
                    "size": 10
                }
            }
        }
    }
    print(search_body)
    resp = es.search(index=target, body=search_body)
    return resp['aggregations']['my_agg']['buckets']
